import argparse
import os
import numpy as np
# import math
# import itertools
import time
import datetime
import sys


import paddle.vision.transforms as transforms
from paddle.vision.transforms import Compose, Resize, Normalize
from utils.save_img import save_image
from paddle.io import DataLoader

# 自定义模块
from models.model import *
# from models.CasNet import *
# from models.ResNet import * 
# from models.ResNetUnet import *
# from models.Unet import *
from utils.datasets import *
# from models.feature import *
# from models.pix2pix import *
# from models.feature import *

import paddle.nn as nn
import paddle.nn.functional as F
import paddle
from PIL import Image

# from tensorboardX import SummaryWriter
from visualdl import LogWriter
paddle.disable_static()
writer = LogWriter(logdir='./log')
path = r"/home/aistudio/data/new_madic_data"
# path = "/home/sd/lijitao/project/MedGAN/new_data/mr"
# path = r"/home/aistudio/data/data118396/madic_data"
# path = "/home/sd/lijitao/project/gans/PyTorch-GAN-master/implementations/pix2pix/dataset/mr"
parser = argparse.ArgumentParser()
parser.add_argument("--use_gpu", type=bool, default=True, help="If use gpu or not.")
parser.add_argument("--epoch", type=int, default=0, help="epoch to start training from")
parser.add_argument("--n_epochs", type=int, default=10000, help="number of epochs of training")
parser.add_argument("--dataset_name", type=str, default="pet2ct", help="name of the dataset")
parser.add_argument("--batch_size", type=int, default=128, help="size of the batches")
parser.add_argument("--lr", type=float, default=0.0002, help="adam: learning rate")
parser.add_argument("--b1", type=float, default=0.5, help="adam: decay of first order momentum of gradient")
parser.add_argument("--b2", type=float, default=0.999, help="adam: decay of first order momentum of gradient")
parser.add_argument("--decay_epoch", type=int, default=100, help="epoch from which to start lr decay")
parser.add_argument("--n_cpu", type=int, default=2, help="number of cpu threads to use during batch generation")
parser.add_argument("--img_height", type=int, default=256, help="size of image height")
parser.add_argument("--img_width", type=int, default=256, help="size of image width")
parser.add_argument("--channels", type=int, default=1, help="number of image channels")
parser.add_argument(
    "--sample_interval", type=int, default=2000, help="interval between sampling of images from generators"
)
parser.add_argument("--checkpoint_interval", type=int, default=50, help="interval between model checkpoints")
opt = parser.parse_args()
# print(opt)
paddle.set_device('gpu:0') if opt.use_gpu else paddle.set_device('cpu')



# paddle.set_device('cpu')

os.makedirs("images/%s" % opt.dataset_name, exist_ok=True)
os.makedirs("saved_models/%s" % opt.dataset_name, exist_ok=True)


criterion_GAN = paddle.nn.MSELoss()
'''L1Loss: 平均绝对误差（MAE）'''
criterion_pixelwise = paddle.nn.L1Loss()

# Loss weight of L1 pixel-wise loss between translated image and real image
# 翻译图像和真实图像
lambda_pixel = 100

# Calculate output of image discriminator (PatchGAN)
patch = (1, opt.img_height // 2 ** 4, opt.img_width // 2 ** 4)

# Initialize generator and discriminator
discriminator = Discriminator(in_channels=3)
generator = GeneratorUNet(in_channels=3, out_channels=3)

if opt.epoch != 0:
    # 加载已经训练好的模型
    generator.load_state_dict(paddle.load("saved_models/%s/generator_%d.pth" % (opt.dataset_name, opt.epoch)))
    discriminator.load_state_dict(paddle.load("saved_models/%s/discriminator_%d.pth" % (opt.dataset_name, opt.epoch)))
else:
    # 初始化权重
    print("paddle init weight")
    # generator.apply(weights_init_normal)
    # discriminator.apply(weights_init_normal)

# 优化器
optimizer_G = paddle.optimizer.Adam(learning_rate = 0.001, parameters = generator.parameters())
optimizer_D = paddle.optimizer.Adam(learning_rate = 0.001, parameters = discriminator.parameters())

# 设置数据加载器

transforms_ = [
            Resize(size = (256, 256)), 
            #    Normalize(mean=[0.5], 
            #              std= [0],
            #              data_format='CHW',
            #              to_rgb=True), 
                ToTensor()]


dataloader = DataLoader(
    ImageDataset(path, transforms_=transforms_, mode='train'),
    batch_size=opt.batch_size,
    shuffle=True,
    num_workers=opt.n_cpu,
    drop_last = True,
)

val_dataloader = DataLoader(
    ImageDataset(path, transforms_=transforms_, mode="val"),
    batch_size=10, 
    shuffle=True,
    num_workers=1,
    drop_last=True,
)

def cal_psnr(im1, im2):
    mse = (np.abs(im1 - im2) ** 2).mean()
    psnr = 10 * np.log10(255 * 255 / mse)
    return psnr


def eval(im1, im2):
    im1 = im1.cpu().detach().numpy()
    im2 = im2.cpu().detach().numpy()
    mse = (np.abs(im1 - im2) ** 2).mean()
    psnr = 10 * np.log(255 * 255 / mse)
    # im1 和 im2 都为灰度图像，uint8 类型
    # print(im1.shape, im2.shape)
    # ssim = skimage.measure.compare_ssim(im1, im2, data_range=255, multichannel=True)
    return mse, psnr    #   , ssim

def mse_psnr(im1, im2):
    im1 = (im1 + 1) / 2
    im2 = (im2 + 1) / 2
    # print(im1.min(), im1.max())
    # print(im2.min(), im2.max())
    mse = paddle.mean((im1/1.0 - im2/1.0) ** 2)
    if mse < 1.0e-10:
        return mse, 100
    return mse, 10 * paddle.log10(1.0**2 / mse)


def sample_images(epoch, i, real_A, real_B, fake_B):
    data, pred, label = real_A * 255, fake_B * 255, real_B * 255
    pred = paddle.clip(pred.detach(), 0, 255)

    data = data.cast('int64')
    pred = pred.cast('int64')
    label = label.cast('int64')
    h, w = pred.shape[-2], pred.shape[-1]
    img = np.zeros((h, 1 * 3 * w, 3))
    for idx in range(0, 1):
        row = idx * h
        tmplist = [data[idx], pred[idx], label[idx]]
        for k in range(3):
            col = k * w
            tmp = np.transpose(tmplist[k], (1, 2, 0))
            img[row:row + h, col:col + w] = np.array(tmp)
    img = img.astype(np.uint8)
    img = Image.fromarray(img)
    if not os.path.exists("./images"):
        os.makedirs("./images")
    img.save("./images/%04d_%06d.png" % (epoch, i))



# ----------
#  Training
# ----------
prev_time = time.time()

for epoch in range(opt.epoch, opt.n_epochs):
    all_loss_G = 0.0
    all_loss_D = 0.0
    all_loss_pixel = 0.0
    all_mse = 0.0
    all_psnr = 0.0

    for i, batch in enumerate(dataloader):

        generator.train()
        discriminator.train()
        # Model inputs
        real_A = batch["A"]
        real_B = batch["B"]
        # print(real_A.cpu().max(), real_B.cpu().max())
        # Adversarial ground truths
        valid = paddle.ones(shape =[opt.batch_size, *patch])
        fake = paddle.zeros(shape = [opt.batch_size, *patch])


        # ------------------
        #  训练生成器
        # ------------------


        # optimizer_G.zero_grad()

        # GAN loss
        # L1 loss

        fake_B = generator(real_A)
        

        # pred_fake = discriminator(fake_B, real_A)
        pred_fake = discriminator(fake_B, real_B)
        
        loss_GAN = criterion_GAN(pred_fake, valid)
        # Pixel-wise loss
        # print(fake_B.min(), fake_B.max())
        loss_pixel = criterion_pixelwise(fake_B, real_B)
        # print(loss_pixel)


        # 特征提取器实现风格迁移损失

        # Style Loss
        # loss_style = feature_loss(fake_B, real_B)

        # Total Loss
        loss_G = loss_GAN + lambda_pixel * loss_pixel   #+ loss_style

        loss_G.backward()
        optimizer_G.step()
        generator.clear_gradients()

        # ---------------------
        #  训练判别器
        # ---------------------

        # optimizer_D.zero_grad()

        # Real loss
        # pred_real = discriminator(real_B, real_A)
        pred_real = discriminator(fake_B.detach(), real_B)

        loss_real = criterion_GAN(pred_real, valid)

        # Fake loss
        pred_fake = discriminator(fake_B.detach(), real_A)
        # pred_fake = discriminator(fake_B.detach(), real_B)
        loss_fake = criterion_GAN(pred_fake, fake)

        # Total loss
        loss_D = 0.5 * (loss_real + loss_fake)

        loss_D.backward()
        # loss_D.backward(retain_graph=True)
        optimizer_D.step()
        discriminator.clear_gradients()

        # --------------
        #  Log Progress
        # --------------

        # 输出训练进度
        batches_done = epoch * len(dataloader) + i
        batches_left = opt.n_epochs * len(dataloader) - batches_done
        time_left = datetime.timedelta(seconds=batches_left * (time.time() - prev_time))
        prev_time = time.time()

        mse, psnr = mse_psnr(real_B.detach(), fake_B.detach())
        # mse = 0.
        # psnr = 0.
        # Print log
        sys.stdout.write(
            "\r[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f, pixel: %f, mse: %f, psnr: %f, adv: %f] ETA: %s"
            % (
                epoch,
                opt.n_epochs,
                i,
                len(dataloader),
                loss_D.item(),
                loss_G.item(),
                loss_pixel.item(),
                mse,
                psnr,
                loss_GAN.item(),
                time_left,
            )
        )
        all_loss_G += loss_G.item()
        all_loss_D += loss_D.item()
        all_loss_pixel += loss_pixel.item()
        all_mse += mse
        all_psnr += psnr

    train_num = i + 1
    # TensorboardX  可视化训练过程
    writer.add_scalar('Loss_G', all_loss_G / train_num, epoch)
    writer.add_scalar('Loss_D', all_loss_D / train_num, epoch)
    writer.add_scalar('loss_pixel', all_loss_pixel / train_num, epoch)
    writer.add_scalar('mse', all_mse / train_num, epoch)
    writer.add_scalar('psnr', all_psnr / train_num, epoch)
        # writer.add_scalar('ssim', ssim, epoch)

    sample_images(epoch, i, real_A, real_B, fake_B)
    # If at sample interval save image
    #     if batches_done % opt.sample_interval == 0:
    #         sample_images(batches_done)

    if opt.checkpoint_interval != -1 and epoch % opt.checkpoint_interval == 0:
        # Save model checkpoints
        paddle.save(generator.state_dict(), "save_models/%s/generator_%d.pth" % (opt.dataset_name, epoch))
        paddle.save(discriminator.state_dict(), "save_models/%s/discriminator_%d.pth" % (opt.dataset_name, epoch))
    
    if epoch % opt.checkpoint_interval == 0:
            current_save_dir = os.path.join("save_models", f'epoch_{epoch}')
            if not os.path.exists(current_save_dir):
                os.makedirs(current_save_dir)
            paddle.save(generator.state_dict(),
                        os.path.join(current_save_dir, 'generator.pdparams'))
            paddle.save(discriminator.state_dict(), 
                        os.path.join(current_save_dir, 'discriminator.pdparams'))
            paddle.save(optimizer_G.state_dict(),
                        os.path.join(current_save_dir, 'generator.pdopt'))
            paddle.save(optimizer_D.state_dict(),
                        os.path.join(current_save_dir, 'discriminator.pdopt'))

writer.close()














